import cx_Oracle, sys

## {{{ Connectionmanager
class ConnectionManager :
	'''handles connection and cursor generation. all the client needs to do is call fetchall'''
	def __init__(s, authstr, authstr_alt, testsql) :
		s.authstr, s.authstr_alt = authstr, authstr_alt
		s.testsql = testsql
		s.conn, s.cursor = None, None
		s.makeCursor()
	def makeCursor(s) :
		try :
			if s.cursor == None or s.conn == None :
				s.conn = cx_Oracle.connect(s.authstr)
				s.cursor = s.conn.cursor()
			s.fetchall(s.testsql)
		except cx_Oracle.Error, e : ## any database error
			print "ConnectionManager using alternate db 1"
			s.conn = cx_Oracle.connect(s.authstr_alt)
			s.cursor = s.conn.cursor()
	def commit(s) :
		s.conn.commit()
	def fetchall(s, query, dontfetch=None) :
		try :
			s.cursor.execute(query)
			if dontfetch != None : return
			return s.cursor.fetchall()
		except cx_Oracle.Error, e : ## any database error
			print "ConnectionManager using alternate db 2"
			s.conn = cx_Oracle.connect(s.authstr_alt)
			s.cursor = s.conn.cursor()
			s.cursor.execute(query)
			if dontfetch != None : return
			return s.cursor.fetchall()
pdbedbauth = 'search_prints/search_prints55@pdbesch6' ; printsdbauth = pdbedbauth 
pdbedbauth_alt = 'search_prints/search_prints55@pdbesch6' ; printsdbauth_alt = pdbedbauth_alt
#pdbedbauth = 'prints/prints55@pdbesch4' ; printsdbauth = pdbedbauth 
#pdbedbauth_alt = 'prints/prints55@pdbesch4' ; printsdbauth_alt = pdbedbauth_alt
#pdbedb = ConnectionManager(pdbedbauth, pdbedbauth_alt, "select id from ENTRY where rownum=1")

pdbedbsch4 = ConnectionManager('prints/prints55@pdbesch4', 'prints/prints55@pdbesch4', "select id from ENTRY where rownum=1")

pdbedb = ConnectionManager('pdberead/pdberead55@pdbesch6', 'pdberead/pdberead55@pdbesch6', "select id from ENTRY where rownum=1")
import atexit
def finalCleanup() :
	pdbedb.conn.close()
	pdbedbsch4.conn.close()
atexit.register(finalCleanup)
## }}}

class EntityInstance :
	'''instance of entity, identifiable by entity-id and struct-asym-id'''
	def __init__(s,entity,said) :
		s.entity, s.said = entity, said
	def getPDBchainID(s) :
		'''get equivalent chain id in PDB file'''
		if not hasattr(s, 'PDBchainID') :
			for row in pdbedb.fetchall("select auth_asym_id from struct_asym where entry_id='%s' and entity_id='%s' and id='%s'" % (s.entity.getEntry().getID(), s.entity.getID(), s.said) ) :
				s.PDBchainID = row[0] ; break
		return s.PDBchainID
	def getDBrowsResidueTableForPolymerOrdered(s) :
		'''gets resnum, icode, resname (in PDB style) from residue table for polymers ordered by their appearance in the polymer chain'''
		sql = '''select res.auth_seq_id, res.pdb_ins_code, res.chem_comp_id
			from entity_poly_seq eps, residue res
			where res.entry_id='%s' and res.struct_asym_id='%s' and res.entity_id='%s'
				and res.entry_id=eps.entry_id and res.entity_id=eps.entity_id and eps.id=res.id
			order by eps.id
		''' % (s.entity.getEntry().getID(), s.said, s.entity.getID())
		return [["%4d"%row[0], ' ' if row[1]==None else row[1], "%3s"%row[2]] for row in pdbedb.fetchall(sql)]
	def setFirstLastResnums(s) :
		'''sets internal fields with residue details in PDB style for first and last observed residues of this polymer'''
		rows = s.getDBrowsResidueTableForPolymerOrdered()
		s.firstPDBresnum = rows[0]
		s.lastPDBresnum  = rows[len(rows)-1]
	def getPDBresnumStart(s) :
		'''first residue in chain actually modelled with atoms (ignoring microhet complication)'''
		if not hasattr(s,'firstPDBres') : s.setFirstLastResnums()
		return s.firstPDBres
	def getPDBresnumEnd(s) :
		'''last residue in chain actually modelled with atoms (ignoring microhet complication)'''
		if not hasattr(s,'firstPDBres') : s.setFirstLastResnums()
		return s.lastPDBres
	def getPDBresnums(s) :
		'''pdb resnum and icode and resname are returned in ordered manner for polymers. does not work for non-polymers yet...'''
		return s.getDBrowsResidueTableForPolymerOrdered()
class Entity :
	def __init__(s,entry,eid) :
		s.entry, s.eid = entry, eid
		s.dbconn = s.entry.dbconn
	def getEntry(s) :
		return s.entry
	def getID(s) :
		return s.eid
	def getInstances(s) :
		if not hasattr(s, 'instances') :
			s.instances = []
			for row in pdbedb.fetchall("select id from struct_asym where entity_id='%s' and entry_id='%s'" % (s.eid, s.entry.getID()) ) :
				s.instances.append(EntityInstance(s,row[0]))
		return s.instances
	def getType(s) :
		if not hasattr(s, 'biochemtype') :
			for row in pdbedb.fetchall("select type, polymer_type from struct_asym where entity_id='%s' and entry_id='%s'" % (s.eid, s.entry.getID()) ) :
				s.biochemtype, s.polytype = row
		return s.biochemtype
	def isPolymer(s) :
		if s.getType() == 'p' : return True
		return False
	def isWater(s) :
		if s.getType() == 'w' : return True
		return False
	def isLigand(s) :
		if s.getType() == 'b' : return True
		return False
	def isSugar(s) :
		if s.getType() == 'p' and s.polytype == 'S' : return True
		return False
	def isDNA(s) :
		if s.getType() == 'p' and s.polytype == 'D' : return True
		return False
	def isRNA(s) :
		if s.getType() == 'p' and s.polytype == 'R' : return True
		return False
	def isProtein(s) :
		if s.getType() == 'p' and s.polytype == 'P' : return True
		return False
	def getDomains(s, domtype) :
		'''return a list of DomMapping objects which describe how scop domains lie on entity's occurences'''
		if not hasattr(s,'doms') : s.doms = {}
		if not s.doms.has_key(domtype) :
			sql, entryid, entid = None, s.getEntry().getID(), s.getID()
			if domtype in ['SCOP','CATH'] :
				domIDfld = {"SCOP":"sccs", "CATH":"cathcode",}
				instIDfld = {"SCOP":'sunid', "CATH":"domain",}
				tablename = {"SCOP":'scop_description', "CATH":"cath_domain",}
				entryid, entid = s.getEntry().getID(), s.getID()
				sql = '''select b.xref_accession, b.pdb_start, b.pdb_end, c.%s, a.entity_id, a.id
					from struct_asym a , xref_region b, %s c
					where b.xref_source = '%s'
					and a.id = b.struct_asym_id
					and a.entry_id = b.entry_id and b.pdb_start is not null and b.pdb_end is not null
					and a.entry_id='%s' and c.%s=b.xref_accession and a.entity_id='%s'
				''' % (domIDfld[domtype], tablename[domtype], domtype, entryid, instIDfld[domtype], entid)
			elif domtype == 'PFAM' :
				sql = '''select dbms_random.random, b.pdb_start, b.pdb_end, b.xref_accession, a.entity_id, a.id
					from struct_asym a , xref_region b
					where b.xref_source = 'PFAM'
					and a.id = b.struct_asym_id
					and a.entry_id = b.entry_id and b.pdb_start is not null and b.pdb_end is not null
					and a.entry_id='%s' and a.entity_id='%s'
				''' % (entryid, entid)
			s.doms[domtype] = []
			domsegs = {}
			for row in pdbedb.fetchall(sql) :
				domInstId, regStart, regEnd, scopSCCS, entity_id, struct_asym_id = row
				if not domsegs.has_key(scopSCCS) : domsegs[scopSCCS] = {}
				if not domsegs[scopSCCS].has_key(domInstId) : domsegs[scopSCCS][domInstId] = []
				domsegs[scopSCCS][domInstId].append((struct_asym_id,regStart,regEnd))
			for scopid in domsegs.keys() :
				for inst in domsegs[scopid] :
					thesaid = None
					for said,start,end in domsegs[scopid][inst] :
						if thesaid == None : thesaid = said
						else : thesaid == said
					s.doms[domtype].append( DomMapping(domtype,scopid,s.getEntry(),s,domsegs[scopid][inst]) )
		if not s.doms.has_key(domtype) : print "ERROR Domains of type", domtype, "not available!" ; sys.exit(1)
		return s.doms[domtype]

class DomMapping :
	'''holds a reference ScopDomain object, and mapping between the domain and an entry-entity-[(structAsymId-start-stop),...].
	refs to full entry entity objects are held. multiple segments are held if required.'''
	def __init__(s, domtype, domid, entry, entity, theranges) :
		s.domtype, s.domid, s.entry, s.entity, s.segments = domtype, domid, entry, entity, theranges
	def getPDBresidueRanges(s) :
		'''return residue ranges in PDB coordinate'''
		pdbsegs = []
		for said,start,stop in s.segments :
			pdbsegs.append( [s.entry.getPDBresidueAddress(s.entity.getID(), said, start)] )
			pdbsegs[len(pdbsegs)-1].append(s.entry.getPDBresidueAddress(s.entity.getID(), said, stop))
		return pdbsegs
	def getID(s) : return s.domid
	def getType(s) : return s.domtype

class PDBEntry :
	def __init__(s, pdbid, dbconn) :
		s.pdbid, s.dbconn = pdbid, dbconn
	def getID(s) :
		return s.pdbid
	def getEntities(s) :
		if not hasattr(s, 'entities') :
			s.entities = []
			for row in pdbedb.fetchall("select id from entity where entry_id='%s'" % s.pdbid) :
				s.entities.append(Entity(s,row[0]))
		return s.entities
	def getEntity(s, eid) :
		for ent in s.getEntities() :
			if eid == ent.getID() : return ent
		return None
	def getDomains(s, domtype) :
		'''returns mapping from entity id to domains of given type, None is none is defined.'''
		if not domtype in ["PFAM","SCOP","CATH"] :
			print "ERROR domain type is unknown:", domtype ; sys.exit(1)
		eid2scop = {}
		for ent in s.getEntities() :
			scopdoms = ent.getDomains(domtype)
			if len(scopdoms) > 0 : eid2scop[ent.getID()] = scopdoms
		return eid2scop
	def getPDBresidueAddress(s, entid, said, seqnum) :
		'''Given entity id, struct asym id, seq num: return chain (1 char), resnum (4 char) and insert-code (1 char) and resname (3 char)'''
		sql = '''select auth_asym_id, auth_seq_id, pdb_ins_code
				from residue
				where entry_id='%s' and entity_id='%s' and struct_asym_id='%s' and id=%d
		''' % ( s.getID(), entid, said, seqnum )
		for chain, resnum, icode in pdbedb.fetchall(sql) :
			if icode == None : icode = ' '
			return chain, '%4s'%resnum, icode
	def getPolymericPDBreslistInPDBchain(s,chid,microhet="ignore") :
		'''list residues in PDB style for polymeric entities in given PDB chain id.
			If there is more than one record for same resnum+icode, only 1 is kept if microhet = ignore. Else all are returned.'''
		retlist = []
		for ent in s.getEntities() :
			if not ent.isPolymer() : continue
			for einst in ent.getInstances() :
				if chid != einst.getPDBchainID() : continue
				resnums = einst.getPDBresnums()
				if microhet == "ignore" : ## most probably this step is unnecessary because residue table doe not contain altcoded or microhet-residues
					resins2ri = {}
					for ri in range(len(resnums)) :
						resnum, icode, resname = resnums[ri]
						if not resins2ri.has_key(resnum+icode) : resins2ri[resnum+icode] = []
						resins2ri[resnum+icode].append(ri)
					markfordeletion = []
					for multris in resins2ri.values() :
						if len(multris) > 1 : markfordeletion += multris[1:]
					markfordeletion.sort(reverse=True)
					for ri in markfordeletion : resnums.pop(ri)
				retlist.append( resnums )
		if len(retlist) > 1 : # order the subists and merge into one list, not needed most of the times
			for ri in range(len(retlist)) :
				for rk in range(ri+1,len(retlist)) :
					sum_i, sum_k = 0, 0
					for rnum,ic,rname in retlist[ri] :
						try : sum_i += string.atoi(rnum)
						except : pass
					for rnum,ic,rname in retlist[rk] :
						try : sum_k += string.atoi(rnum)
						except : pass
					if sum_i / len(retlist[ri]) > sum_k / len(retlist[rk]) :
						tlist = retlist[ri]
						retlist[ri] = retlist[rk]
						retlist[rk] = tlist
			for ri in range(1,len(retlist)) : retlist[0] += retlist[ri]
		if len(retlist) > 0 : retlist = [x for x in retlist[0]]
		return retlist
	def getDomainsForPDBchain(s,chid, domtype) :
		'''returns how scop domains map onto given PDB chain id. returns {id1:[ranges], id2:[ranges],...}. ranges are in PDB style.'''
		ret = {}
		for eid,dommaps in s.getDomains(domtype).items() :
			for dmap in dommaps :
				domid, ranges = dmap.getID(), dmap.getPDBresidueRanges()
				samechain = []
				for arange in ranges :
					if arange[0][0] != chid or arange[1][0] != chid : continue
					else : samechain.append(arange)
				if len(samechain) == 0 : continue
				if not ret.has_key(domid) : ret[domid] = []
				ret[domid].append(samechain)
		return ret

if __name__ == "__main__" :
	pe = PDBEntry('4dow', pdbedb)
	for eid,dommaps in pe.getDomains('SCOP').items() :
		for dmap in dommaps :
			print dmap.getID(), dmap.getPDBresidueRanges()
	for ent in pe.getEntities() :
		#ent = pe.getEntity(eid)
		if not ent.isProtein() : continue
		for einst in ent.getInstances() :
			for pdbresdetails in einst.getPDBresnums() :
				pass #print pdbresdetails
	#print pe.getPolymericPDBreslistInPDBchain('A')
	print pe.getDomainsForPDBchain('A','SCOP')
	print pe.getDomainsForPDBchain('A','CATH')
	print pe.getDomainsForPDBchain('A','PFAM')
